import cv2
import numpy as np
import onnxruntime as ort

# Load MNIST ONNX model
session = ort.InferenceSession("mnist-8.onnx", providers=['CPUExecutionProvider'])
input_name = session.get_inputs()[0].name

# Digit recognition function
def recognize_digit(roi):
    roi_resized = cv2.resize(roi, (28, 28))
    roi_normalized = roi_resized.astype(np.float32) / 255.0
    roi_input = roi_normalized.reshape(1, 1, 28, 28)
    output = session.run(None, {input_name: roi_input})[0]
    pred = np.argmax(output)
    conf = float(np.max(output))
    return pred, conf

# Open camera
cap = cv2.VideoCapture(0)
assert cap.isOpened(), "Failed to open camera"

while True:
    ret, frame = cap.read()
    if not ret:
        break

    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    _, thresh = cv2.threshold(blurred, 128, 255, cv2.THRESH_BINARY_INV)

    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if w < 10 or h < 10:
            continue
        digit_roi = thresh[y:y+h, x:x+w]
        pred, conf = recognize_digit(digit_roi)
        label = f"{pred} ({conf:.2f})"
        cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
        cv2.putText(frame, label, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (0, 0, 255), 1)

    cv2.imshow("Digit Detection", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()